import cv2
import numpy as np

def detect_defects(reference_img_path, test_img_path, output_path):
    # 1. 读取图片并统一尺寸
    ref_img = cv2.imread(reference_img_path)
    test_img = cv2.imread(test_img_path)
    if ref_img is None or test_img is None:
        raise ValueError("无法读取图片，请检查路径")
    
    # 统一尺寸（确保两张图大小一致）
    ref_img = cv2.resize(ref_img, (640, 480))
    test_img = cv2.resize(test_img, (640, 480))
    
    # 2. 预处理：转为灰度图并高斯模糊（降噪）
    ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY)
    test_gray = cv2.cvtColor(test_img, cv2.COLOR_BGR2GRAY)
    
    ref_blur = cv2.GaussianBlur(ref_gray, (5, 5), 0)
    test_blur = cv2.GaussianBlur(test_gray, (5, 5), 0)
    
    # 3. 计算差异并阈值化
    diff = cv2.absdiff(ref_blur, test_blur)  # 绝对值差
    _, thresh = cv2.threshold(diff, 30, 255, cv2.THRESH_BINARY)  # 阈值化（30为经验值，可调整）
    
    # 4. 形态学操作：去除噪点，强化缺陷区域
    kernel = np.ones((3, 3), np.uint8)
    eroded = cv2.erode(thresh, kernel, iterations=1)  # 腐蚀（去除小噪点）
    dilated = cv2.dilate(eroded, kernel, iterations=2)  # 膨胀（恢复缺陷大小）
    
    # 5. 检测轮廓并标注缺陷
    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    for contour in contours:
        # 过滤过小的轮廓（避免误检）
        area = cv2.contourArea(contour)
        if area < 50:  # 面积阈值可根据实际场景调整
            continue
        
        # 获取缺陷边界框
        x, y, w, h = cv2.boundingRect(contour)
        # 用红色矩形标注缺陷（(0,0,255)为BGR格式的红色）
        cv2.rectangle(test_img, (x, y), (x + w, y + h), (0, 0, 255), 2)
        # 可添加文字标签
        cv2.putText(test_img, "Defect", (x, y - 10), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    
    # 保存结果
    cv2.imwrite(output_path, test_img)
    print(f"缺陷标注结果已保存至：{output_path}")
    
    # 显示结果（可选）
    cv2.imshow("Defect Detection", test_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# 示例调用
if __name__ == "__main__":
    reference_path = "reference.jpg"  # 标准无缺陷图
    test_path = "test.jpg"            # 待检测图
    output_path = "result.jpg"        # 标注结果图
    detect_defects(reference_path, test_path, output_path)